#!/usr/bin/python
# -*- coding: UTF-8 -*-
import os

import numpy as np
import torch.nn as nn
import torch
import torch.nn.functional as F
from model_service.pytorch_model_service import PTServingBaseService
# from models.dnns_residual import DNNs
from models.dnns import DNNs

# class Mish(nn.Module):
#     def __init__(self):
#         super().__init__()
#         print("Mish activation loaded...")
#     def forward(self,x):
#         x = x * (torch.tanh(F.softplus(x)))
#         return x
# class DNN(nn.Module):
#     def __init__(self, in_dim,
#                  n_hidden_1, n_hidden_2,
#                  num_class, dropout_p=0.5):
#         super().__init__()
#         self.layer1 = nn.Linear(in_dim, n_hidden_1)
#         self.layer2 = nn.Linear(n_hidden_1, n_hidden_2)
#         self.layer3 = nn.Linear(n_hidden_2, n_hidden_2)
#         self.layer4= nn.Linear(n_hidden_2, n_hidden_2)
#         self.layer5 = nn.Linear(n_hidden_2, num_class)
#         self.activator = Mish()
#         self.dropout0 = nn.Dropout(p=0.3)
#         self.dropout1 = nn.Dropout(p=dropout_p)
#         self.dropout2 = nn.Dropout(p=dropout_p)
#         self.dropout3 = nn.Dropout(p=dropout_p)
#         self.dropout4 = nn.Dropout(p=dropout_p)
#         # self.softmax = nn.Softmax(dim=1)
#
#     def forward(self, x):
#         x = self.dropout0(x)
#         x = self.layer1(x)
#         x = self.activator(x)
#         x = self.dropout1(x)
#         x = self.layer2(x)
#         x = self.activator(x)
#         x = self.dropout2(x)
#         x = self.layer3(x)
#         x = self.dropout3(x)
#         x = self.layer4(x)
#         x = self.dropout4(x)
#         x = self.layer5(x)
#         # x = self.softmax(x)
#         return x
# class DNNs(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.dnn0 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
#         self.dnn1 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
#         self.dnn2 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
#         self.dnn3 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
#         self.dnn4 = DNN(28169, 4096, 512, 3, dropout_p=0.4)
#         # for i,dnn in enumerate([self.dnn0,self.dnn1,self.dnn2,self.dnn3,self.dnn4]):
#         #     pre_weight = torch.load(pth_path_list[i])
#         #     dnn.load_state_dict(pre_weight)
#
#
#     def forward(self, x):
#         output0=self.dnn0(x)
#         output1=self.dnn1(x)
#         output2=self.dnn2(x)
#         output3=self.dnn3(x)
#         output4=self.dnn4(x)
#         output = output0+output1+output2+output3+output4
#         output = output/5
#
#         return output
import xgboost as xgb


class PredictService(PTServingBaseService):
    def __init__(self, model_name, model_path):
        super(PredictService, self).__init__(model_name, model_path)

        print('-----------------Starting Done----------------')
        dir_path = os.path.dirname(os.path.realpath(model_path))
        xgb_path = os.path.join(dir_path, 'xgboost.json')
        xgb_path_0 = os.path.join(dir_path, 'xgboost_0.json')
        xgb_path_1 = os.path.join(dir_path, 'xgboost_1.json')
        xgb_path_2 = os.path.join(dir_path, 'xgboost_2.json')
        xgb_path_3 = os.path.join(dir_path, 'xgboost_3.json')
        xgb_path_4 = os.path.join(dir_path, 'xgboost_4.json')
        # lgbm_path = os.path.join(dir_path, 'lgbm.txt')
        self.xgb = xgb.XGBClassifier()
        self.xgb.load_model(xgb_path)
        self.xgb_0 = xgb.XGBClassifier()
        self.xgb_0.load_model(xgb_path_0)
        self.xgb_1 = xgb.XGBClassifier()
        self.xgb_1.load_model(xgb_path_1)
        self.xgb_2 = xgb.XGBClassifier()
        self.xgb_2.load_model(xgb_path_2)
        self.xgb_3 = xgb.XGBClassifier()
        self.xgb_3.load_model(xgb_path_3)
        self.xgb_4 = xgb.XGBClassifier()
        self.xgb_4.load_model(xgb_path_4)


        model = DNNs()
        state_dict = torch.load(model_path)
        model.load_state_dict(state_dict)
        model.eval()
        self.model = model
        self.load_preprocess()

    def load_preprocess(self, mean_name='mean.npy', std_name='std.npy'):
        dir_path = os.path.dirname(os.path.realpath(self.model_path))

        mean_path = os.path.join(dir_path, mean_name)
        std_path = os.path.join(dir_path, std_name)
        self.mean = np.load(mean_path)
        self.std = np.load(std_path)

    def _preprocess(self, data):
        print('pre_data:{}'.format(data))
        preprocessed_data = {}

        for d in data:
            for k, v in data.items():
                for file_name, features_path in v.items():
                    x = np.load(features_path)
                    # deploy environment numpy version
                    x = (x - self.mean) / self.std
                    x = np.nan_to_num(x)
                    x[x > 1000000] = 0
                    x[x < -1000000] = 0
                    x = torch.from_numpy(x).to(torch.float32)
                    preprocessed_data[k] = x
        print('preprocessed_data', preprocessed_data)
        return preprocessed_data

    def _postprocess(self, data):
        print('post_data:{}'.format(data))
        infer_output = {}

        for output_name, result in data.items():
            infer_output['scores'] = result.tolist()
        print('infer_output', infer_output)
        return infer_output

    def _inference(self, data):
        print('----infer:-----', data)
        result = {}
        for k, v in data.items():
            v = v.view(1, -1)
            output, x1_0, x1_1, x1_2, x1_3, x1_4 = self.model(v)
            print('output:', output, type(output))
            print('x1_0:', x1_0, type(x1_0))
            x1 = v.detach().cpu().numpy()
            x1_0 = x1_0.detach().cpu().numpy()
            x1_1 = x1_1.detach().cpu().numpy()
            x1_2 = x1_2.detach().cpu().numpy()
            x1_3 = x1_3.detach().cpu().numpy()
            x1_4 = x1_4.detach().cpu().numpy()
            output_xgb = self.xgb.predict_proba(x1)
            output_0 = self.xgb_0.predict_proba(x1_0)
            output_1 = self.xgb_1.predict_proba(x1_1)
            output_2 = self.xgb_2.predict_proba(x1_2)
            output_3 = self.xgb_3.predict_proba(x1_3)
            output_4 = self.xgb_4.predict_proba(x1_4)
            output_catb = self.catboost.predict_proba(x1)
            # output_lgbm = self.lgbm.predict(x1)
            output = (output_xgb + output_catb + output_0 + output_1 + output_2 + output_3 + output_4 + output.detach().cpu().numpy() * 6) / 13
            result[k] = output
        print('----infer_result----:', result)
        return result
